import numpy as np
import torch
from typing import List, Dict, Optional, Tuple, Union
from mmengine.utils import is_list_of
from pprint import pprint
import math
from mmengine.optim import build_optim_wrapper
from train_utils import build_param_scheduler
from torch.optim.sgd import SGD
from mmengine.optim.optimizer.optimizer_wrapper import OptimWrapper


def mmengine_collate(input_batch):
    images = []
    data_samples = []
    for sample in input_batch:
        images.append(sample['inputs'])
        data_samples.append(sample['data_samples'])

    return {'inputs':images, 'data_samples':data_samples}

def set_configs(size):
    # ins
    crop_size = (size, size)
    batch_augments = [
        dict(
            type='BatchFixedSizePad',
            size=crop_size,
            img_pad_value=0,
            pad_mask=True,
            mask_pad_value=0,
            pad_seg=False,
            seg_pad_value=255)
    ]

    train_pipeline = [
        dict(type='MTP_LoadImageFromFile', to_float32=True),
        dict(type='MTP_LoadAnnotations', with_bbox=True, with_rbox=True, with_cls = True,
            with_mask=True, with_seg=True, box_type = 'hbox', rbox_type='qbox',
            reduce_zero_label=False, ignore_index = 255),
        dict(type='ROD_ConvertBoxType', box_type_mapping=dict(gt_rboxes='rbox')),
        dict(
            type='MTP_RandomFlip',
            prob=0.75,
            direction=['horizontal', 'vertical', 'diagonal']),
        dict(
            type='MTP_RandomResize',
            scale=crop_size,
            ratio_range=(0.5, 2.0),
            resize_type='MTP_Resize',
            keep_ratio=True),
        dict(
            type='MTP_RandomCrop',
            crop_size=crop_size,
            crop_type='absolute',
            recompute_bbox=True,
            allow_negative_crop=True),
 
        # dict(type='INS_FilterAnnotations', min_gt_bbox_wh=(1e-5, 1e-5), by_mask=True),
        dict(type='MTP_PhotoMetricDistortion'),
        dict(type='MTP_PackInputs')
    ]

    valid_pipeline = [
        dict(type='MTP_LoadImageFromFile', to_float32=True),
        dict(type='MTP_Resize', scale=crop_size, keep_ratio=True),
        # avoid bboxes being resized
        dict(type='MTP_LoadAnnotations', with_bbox=True, with_rbox=True, with_cls = True,
            with_mask=True, with_seg=True, box_type = 'hbox', rbox_type='qbox',
            reduce_zero_label=False, ignore_index = 255),
        dict(type='ROD_ConvertBoxType', box_type_mapping=dict(gt_rboxes='rbox')),
        dict(
            type='MTP_PackInputs',
            meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', 'scale_factor'))
    ]

    return train_pipeline, valid_pipeline, batch_augments


def parse_datainfos(dataset, img_ids):

    data_list = []
    total_ann_ids = []
    
    for img_id in img_ids:
        img_id = int(img_id)
        raw_img_info = dataset.coco.load_imgs([img_id])[0] # 图片信息
        # pprint(raw_img_info)
        raw_img_info['img_id'] = img_id

        ann_ids = dataset.coco.get_ann_ids(img_ids=[img_id]) # 得到每张图包含的ann id
        raw_ann_info = dataset.coco.load_anns(ann_ids) # 这些ann id 对应的ann信息
        total_ann_ids.extend(ann_ids)

        parsed_data_info = dataset.parse_data_info({
            'raw_ann_info': raw_ann_info,
            'raw_img_info': raw_img_info
        })
        data_list.append(parsed_data_info)
    if dataset.ANN_ID_UNIQUE: # ann id必须是独一无二的
        assert len(set(total_ann_ids)) == len(
            total_ann_ids
        ), f"Annotation ids in '{dataset.ann_file}' are not unique!"

    #del self.coco
    
    return data_list

def data_augs(dataset, datainfos):

    data_samples = []

    for i in range(len(datainfos)):

        data_info = datainfos[i]

        data_sample = dataset.pipeline(data_info)

        while data_sample is None:

            idx = np.random.randint(0, dataset.length)

            raw_img_info = dataset.coco.load_imgs([dataset.img_ids[idx]])[0] # 图片信息
            raw_img_info['img_id'] = dataset.img_ids[idx]

            ann_ids = dataset.coco.get_ann_ids(img_ids=[dataset.img_ids[idx]]) # 得到每张图包含的ann id
            raw_ann_info = dataset.coco.load_anns(ann_ids) # 这些ann id 对应的ann信息

            parsed_data_info = dataset.parse_data_info({
                'raw_ann_info': raw_ann_info,
                'raw_img_info': raw_img_info
            })

            data_sample = dataset.pipeline(parsed_data_info)

        data_samples.append(data_sample)
    data_samples = mmengine_collate(data_samples)

    return data_samples


def parse_losses(losses: List[Dict[str, torch.Tensor]]):
    """Parses the raw outputs (losses) of the network.

    Args:
        losses (dict): Raw output of the network, which usually contain
            losses and other necessary information.

    Returns:
        tuple[Tensor, dict]: There are two elements. The first is the
        loss tensor passed to optim_wrapper which may be a weighted sum
        of all losses, and the second is log_vars which will be sent to
        the logger.
    """
    # pprint(losses)
    loss_sum = 0
    # loss_datasets = []

    log_vars = []
    for loss_name, loss_value in losses.items():
        if isinstance(loss_value, torch.Tensor):
            log_vars.append([loss_name, loss_value.mean()])
        elif is_list_of(loss_value, torch.Tensor):
            log_vars.append(
                [loss_name, sum(_loss.mean() for _loss in loss_value)])
        else:
            raise TypeError(
                f'{loss_name} is not a tensor or list of tensors')
    # print(log_vars)
    loss = sum(value for key, value in log_vars if 'loss' in key)

    loss_sum += loss

    # loss_datasets.append(loss)

    return loss_sum


def get_optimizer(args, model, train_loader):
        # get optimizer
    if 'vit_' in args.backbone:
        # AdamW optimizer, no weight decay for position embedding & layer norm in backbone
        if 'vit_b' in args.backbone:
            optim_wrapper = dict(
                optimizer=dict(
                type='AdamW', lr=1e-4, betas=(0.9, 0.999), weight_decay=0.05),
                constructor='LayerDecayOptimizerConstructor_ViT', 
                paramwise_cfg=dict(
                    num_layers=12, 
                    layer_decay_rate=0.75,
                    )
                    )  
        elif 'vit_l' in args.backbone:
            optim_wrapper = dict(
                optimizer=dict(
                type='AdamW', lr=1e-4, betas=(0.9, 0.999), weight_decay=0.05),
                constructor='LayerDecayOptimizerConstructor_ViT', 
                paramwise_cfg=dict(
                    num_layers=24, 
                    layer_decay_rate=0.9,
                    )
                    )
            
        else:
            raise NotImplementedError
        param_scheduler = [
            dict(
                type='LinearLR',
                start_factor=1.0 / 3,
                by_epoch=False,
                begin=0,
                end=500), 
            dict(
                type='MultiStepLR',
                begin=0,
                end=args.max_epoch,
                by_epoch=True,
                milestones=[16, 22],
                gamma=0.1)
        ]
        optimizer = build_optim_wrapper(model, optim_wrapper)
        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer.optimizer, args.max_epoch*len(train_dataset), eta_min=0, last_epoch=-1)
        scheduler = build_param_scheduler(optimizer, param_scheduler, train_loader)

    elif 'internimage' in args.backbone:
        # classification & segmentation
        optim_wrapper = dict(
            optimizer=dict(
            type='AdamW', lr=2e-5, betas=(0.9, 0.999), weight_decay=0.05),
            constructor='CustomLayerDecayOptimizerConstructor_InternImage',
            paramwise_cfg=dict(num_layers=39, 
                            layer_decay_rate=0.94,
                            depths=[5, 5, 24, 5]
                            )
                            )
        optimizer = build_optim_wrapper(model, optim_wrapper)
        param_scheduler = [
            dict(
                type='LinearLR',
                start_factor=1.0 / 3,
                by_epoch=False,
                begin=0,
                end=500), 
            dict(
                type='MultiStepLR',
                begin=0,
                end=args.max_epoch,
                by_epoch=True,
                milestones=[16, 22],
                gamma=0.1)
        ]
        scheduler = build_param_scheduler(optimizer, param_scheduler, train_loader)
        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer.optimizer, args.max_epoch*len(train_dataset), eta_min=0, last_epoch=-1)
    elif 'vitae' in args.backbone:
        optim_wrapper = dict(
            optimizer=dict(
            type='AdamW', lr=1e-4, betas=(0.9, 0.999),  weight_decay=0.05),
            # constructor='CustomLayerDecayOptimizerConstructor',
            paramwise_cfg=dict(
            custom_keys={'absolute_pos_embed': dict(decay_mult=0.),
                             'relative_position_bias_table': dict(decay_mult=0.),
                             'norm': dict(decay_mult=0.)}
                             )                     
        )
        # optimizer = torch.optim.AdamW()
        optimizer = build_optim_wrapper(model, optim_wrapper)
        param_scheduler = [
            dict(
                type='LinearLR',
                start_factor=1.0 / 3,
                by_epoch=False,
                begin=0,
                end=1500), 
            dict(
                type='MultiStepLR',
                begin=0,
                end=args.max_epoch,
                by_epoch=True,
                milestones=[16, 22],
                gamma=0.1)
        ]
        scheduler = build_param_scheduler(optimizer, param_scheduler, train_loader)
        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer.optimizer, args.max_epoch*len(train_dataset), eta_min=0, last_epoch=-1)
    elif 'resnet' in args.backbone:
        optim_wrapper = dict(
            type=OptimWrapper,
            optimizer=dict(type=SGD, lr=5e-3, momentum=0.9, weight_decay=0.0001))
        param_scheduler = [
            dict(
                type='LinearLR',
                start_factor=1.0 / 3,
                by_epoch=False,
                begin=0,
                end=500), 
            dict(
                type='MultiStepLR',
                begin=0,
                end=args.max_epoch,
                by_epoch=True,
                milestones=[16, 22],
                gamma=0.1)
        ]
        optimizer = build_optim_wrapper(model, optim_wrapper)
        scheduler = build_param_scheduler(optimizer, param_scheduler, train_loader)

    elif 'swin' in args.backbone:
        optim_wrapper = dict(
            optimizer=dict(
            type='AdamW', lr=0.00006, betas=(0.9, 0.999),  weight_decay=0.05),
            # constructor='CustomLayerDecayOptimizerConstructor',
            paramwise_cfg=dict(
            custom_keys={'absolute_pos_embed': dict(decay_mult=0.),
                             'relative_position_bias_table': dict(decay_mult=0.),
                             'norm': dict(decay_mult=0.)}
                             )                     
        )
        param_scheduler = [
            dict(
                type='LinearLR',
                start_factor=1.0 / 3,
                by_epoch=False,
                begin=0,
                end=500), 
            dict(
                type='MultiStepLR',
                begin=0,
                end=args.max_epoch,
                by_epoch=True,
                milestones=[16, 22],
                gamma=0.1)
        ]
        optimizer = build_optim_wrapper(model, optim_wrapper)
        scheduler = build_param_scheduler(optimizer, param_scheduler, train_loader)

    return optimizer, scheduler


class AveragePrecisionMeter(object):
    """
    The APMeter measures the average precision per class.
    The APMeter is designed to operate on `NxK` Tensors `output` and
    `target`, and optionally a `Nx1` Tensor weight where (1) the `output`
    contains model output scores for `N` examples and `K` classes that ought to
    be higher when the model is more convinced that the example should be
    positively labeled, and smaller when the model believes the example should
    be negatively labeled (for instance, the output of a sigmoid function); (2)
    the `target` contains only values 0 (for negative examples) and 1
    (for positive examples); and (3) the `weight` ( > 0) represents weight for
    each sample.
    """

    def __init__(self, difficult_examples=False):
        super(AveragePrecisionMeter, self).__init__()
        self.reset()
        self.difficult_examples = difficult_examples

    def reset(self):
        """Resets the meter with empty member variables"""
        self.scores = torch.FloatTensor(torch.FloatStorage())
        self.targets = torch.LongTensor(torch.LongStorage())

    def add(self, output, target):
        """
        Args:
            output (Tensor): NxK tensor that for each of the N examples
                indicates the probability of the example belonging to each of
                the K classes, according to the model. The probabilities should
                sum to one over all classes
            target (Tensor): binary NxK tensort that encodes which of the K
                classes are associated with the N-th input
                    (eg: a row [0, 1, 0, 1] indicates that the example is
                         associated with classes 2 and 4)
            weight (optional, Tensor): Nx1 tensor representing the weight for
                each example (each weight > 0)
        """
        if not torch.is_tensor(output):
            output = torch.from_numpy(output)
        if not torch.is_tensor(target):
            target = torch.from_numpy(target)

        if output.dim() == 1:
            output = output.view(-1, 1)
        else:
            assert output.dim() == 2, \
                'wrong output size (should be 1D or 2D with one column \
                per class)'
        if target.dim() == 1:
            target = target.view(-1, 1)
        else:
            assert target.dim() == 2, \
                'wrong target size (should be 1D or 2D with one column \
                per class)'
        if self.scores.numel() > 0:
            assert target.size(1) == self.targets.size(1), \
                'dimensions for output should match previously added examples.'

        # make sure storage is of sufficient size
        if self.scores.storage().size() < self.scores.numel() + output.numel():
            new_size = math.ceil(self.scores.storage().size() * 1.5)
            self.scores.storage().resize_(int(new_size + output.numel()))
            self.targets.storage().resize_(int(new_size + output.numel()))
        
        # store scores and targets
        offset = self.scores.size(0) if self.scores.dim() > 0 else 0
        self.scores.resize_(offset + output.size(0), output.size(1))
        self.targets.resize_(offset + target.size(0), target.size(1))
        self.scores.narrow(0, offset, output.size(0)).copy_(output)
        self.targets.narrow(0, offset, target.size(0)).copy_(target)
        # print(self.scores.shape, self.targets.shape)

    def value(self):
        """Returns the model's average precision for each class
        Return:
            ap (FloatTensor): 1xK tensor, with avg precision for each class k
        """

        if self.scores.numel() == 0:
            return 0
        ap = torch.zeros(self.scores.size(1))
        rg = torch.arange(1, self.scores.size(0)).float()
        # compute average precision for each class
        for k in range(self.scores.size(1)):
            # sort scores
            scores = self.scores[:, k]
            targets = self.targets[:, k]
            # compute average precision
            ap[k] = AveragePrecisionMeter.average_precision(scores, targets, self.difficult_examples)
        return ap

    @staticmethod
    def average_precision(output, target, difficult_examples=False):

        # sort examples
        sorted, indices = torch.sort(output, dim=0, descending=True)

        # Computes prec@i
        pos_count = 0.
        total_count = 0.
        precision_at_i = 0.
        for i in indices:
            label = target[i]
            if difficult_examples and label == 0:
                continue
            if label == 1:
                pos_count += 1
            total_count += 1
            if label == 1:
                precision_at_i += pos_count / total_count
        precision_at_i /= pos_count
        return precision_at_i

    def overall(self):
        if self.scores.numel() == 0:
            return 0
        scores = self.scores.cpu().numpy()
        targets = self.targets.cpu().numpy()
        targets[targets == -1] = 0
        return self.evaluation(scores, targets)

    def overall_topk(self, k):
        targets = self.targets.cpu().numpy()
        targets[targets == -1] = 0
        n, c = self.scores.size()
        scores = np.zeros((n, c)) - 1
        index = self.scores.topk(k, 1, True, True)[1].cpu().numpy()
        tmp = self.scores.cpu().numpy()
        for i in range(n):
            for ind in index[i]:
                scores[i, ind] = 1 if tmp[i, ind] >= 0 else -1
        return self.evaluation(scores, targets)


    def evaluation(self, scores_, targets_):
        n, n_class = scores_.shape
        Nc, Np, Ng = np.zeros(n_class), np.zeros(n_class), np.zeros(n_class)
        for k in range(n_class):
            scores = scores_[:, k]
            targets = targets_[:, k]
            targets[targets == -1] = 0
            Ng[k] = np.sum(targets == 1)
            Np[k] = np.sum(scores >= 0)
            Nc[k] = np.sum(targets * (scores >= 0))
        Np[Np == 0] = 1
        OP = np.sum(Nc) / np.sum(Np)
        OR = np.sum(Nc) / np.sum(Ng)
        OF1 = (2 * OP * OR) / (OP + OR)

        CP = np.sum(Nc / Np) / n_class
        CR = np.sum(Nc / Ng) / n_class
        CF1 = (2 * CP * CR) / (CP + CR)
        print(OP, OR, OF1, CP, CR, CF1)
        return OP, OR, OF1, CP, CR, CF1